import tensorflow as tf
from tensorflow_addons.optimizers import MultiOptimizer

import numpy as np


def get_model():
    inputs = tf.keras.Input(shape=(4,))
    outputs = tf.keras.layers.Dense(32)(inputs)
    return tf.keras.Model(inputs, outputs)


model1 = get_model()
model2 = get_model()
model3 = get_model()

inputs = tf.keras.Input(shape=(4,))
y1 = model1(inputs)
y2 = model2(inputs)
y3 = model3(inputs)
outputs = tf.keras.layers.Average()([y1, y2, y3])
model = tf.keras.Model(inputs, outputs)

optimizers_and_layers = [
    (tf.keras.optimizers.SGD(), model1),
    (tf.keras.optimizers.SGD(learning_rate=0.0), model2),
    (tf.keras.optimizers.SGD(), model3),
]



multi_optimizer = MultiOptimizer(optimizers_and_layers)

model.compile(multi_optimizer, loss="mse")

x = np.ones((128, 4)).astype(np.float32)
y = np.ones((128, 32)).astype(np.float32)
model.fit(x, y)
model.save("tf_model",include_optimizer=False)

model1 = tf.keras.models.load_model('tf_model')